import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image
import cv2
import numpy as np
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from Dark_Channel_Prior import *
from Histogram_Equalization import *
from clahe_dehaze import *


def load_image():
    file_path = filedialog.askopenfilename()
    if file_path:
        lbl_image_path.config(text=f"加载的图像: {file_path}")
        img = cv2.imread(file_path)
        lbl_image_path.image_path = file_path
        display_images(img, None)


def dehaze_image():
    if not hasattr(lbl_image_path, "image_path"):
        messagebox.showerror("错误", "请先加载图像")
        return

    img_path = lbl_image_path.image_path
    img = cv2.imread(img_path)

    progress_bar.start()

    # 根据用户选择的去雾算法执行相应的处理
    selected_algorithm = algorithm_combobox.get()

    if selected_algorithm == "DarkChannel":
        dehazed_img = dehaze_algorithm_1(img)
    elif selected_algorithm == "EqualizeHist":
        dehazed_img = dehaze_algorithm_2(img)
    elif selected_algorithm == "CLAHE":
        dehazed_img = dehaze_algorithm_3(img)
    else:
        messagebox.showerror("错误", "请选择有效的去雾算法")
        progress_bar.stop()
        return

    progress_bar.stop()

    # 显示原图和去雾后的图像
    display_images(img, dehazed_img)

def dehaze_algorithm_3(img):
    print("对比度受限的自适应直方图均衡化clahe")
    result = clahe_dehaze(img)

    return result

def dehaze_algorithm_2(img):
    print("直方图均衡化EqualizeHist")

    result = equalize_hist(img)


    return result

def dehaze_algorithm_1(img):
    print("暗通道先验DarkChannel")
    result = ret_dehaze(img)

    return result


def save_image():
    if not hasattr(lbl_image_path, "dehazed_img"):
        messagebox.showerror("错误", "请先处理图像")
        return

    file_path = filedialog.asksaveasfilename(defaultextension=".png",
                                             filetypes=[("PNG files", "*.png"), ("JPEG files", "*.jpg"),
                                                        ("All files", "*.*")])
    if file_path:
        dehazed_img = lbl_image_path.dehazed_img
        if dehazed_img.max() <= 1.0:
            dehazed_img = (dehazed_img * 255).astype(np.uint8)
        cv2.imwrite(file_path, dehazed_img)
        messagebox.showinfo("保存成功", f"图像已保存到: {file_path}")


def display_images(original_img, dehazed_img):
    fig.clear()

    # 显示原始图像
    ax1 = fig.add_subplot(1, 2, 1)
    ax1.imshow(cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB))
    ax1.set_title("source image")
    ax1.axis('off')

    if dehazed_img is not None:
        # 显示去雾后的图像
        ax2 = fig.add_subplot(1, 2, 2)
        #dehazed_img_8bit = (dehazed_img * 255).astype(np.uint8)
        ax2.imshow(cv2.cvtColor(dehazed_img, cv2.COLOR_BGR2RGB))
        ax2.set_title("dehazed image")
        ax2.axis('off')

        # 保存去雾后的图像到标签中
        lbl_image_path.dehazed_img = dehazed_img

    # 更新画布
    canvas.draw()


# 创建主窗口
root = tk.Tk()
root.title("图像去雾程序")
root.geometry("900x600")

# 创建并放置按钮和标签
frame_buttons = tk.Frame(root)
frame_buttons.pack(pady=10)

btn_load = tk.Button(frame_buttons, text="加载图像", command=load_image)
btn_load.grid(row=0, column=0, padx=10)

btn_dehaze = tk.Button(frame_buttons, text="去雾", command=dehaze_image)
btn_dehaze.grid(row=0, column=1, padx=10)

btn_save = tk.Button(frame_buttons, text="保存图像", command=save_image)
btn_save.grid(row=0, column=2, padx=10)

lbl_image_path = tk.Label(root, text="加载的图像: 无")
lbl_image_path.pack(pady=5)

# 算法选择框
algorithm_combobox = ttk.Combobox(root, values=["DarkChannel", "EqualizeHist","CLAHE"])
algorithm_combobox.current(0)  # 默认选择第一个算法
algorithm_combobox.pack(pady=20)


# 进度条
progress_bar = ttk.Progressbar(root, mode='indeterminate')
progress_bar.pack(pady=5)

# 创建 Matplotlib 图形并将其嵌入到 Tkinter 中
fig = Figure(figsize=(8, 4))
canvas = FigureCanvasTkAgg(fig, master=root)
canvas.get_tk_widget().pack(pady=20)

# 启动主循环
root.mainloop()